import numpy as np

# 参数设置
N = 64  # 子载波数量
fftsize = N  # FFT大小
cpsize = 16  # 循环前缀大小
mimo_antennas = 4  # MIMO天线数量
databits_persymbol = 2  # 每个QAM符号的比特数

# 生成随机数据比特
np.random.seed(0)
databits = np.random.randint(0, 2, (N // 2, databits_persymbol))

# 生成QAM符号
def generate_qam_symbols(databits):
    qam_points = np.array([1+1j, 1-1j, -1+1j, -1-1j])  # 4-QAM 星座点
    symbols = qam_points[databits[:, 0] * 2 + databits[:, 1]]
    return symbols

# OFDM调制
def ofdm_modulate(symbols, fftsize, cpsize):
    ofdm_time = np.fft.ifft(np.fft.fftshift(symbols), fftsize)
    ofdm_with_cp = np.concatenate((ofdm_time[-cpsize:], ofdm_time))
    return ofdm_with_cp

# 信道模型
def channel_model(transmitted_signal, h):
    received_signal = np.zeros((len(transmitted_signal), mimo_antennas), dtype=complex)
    for i in range(mimo_antennas):
        received_signal[:, i] = np.convolve(transmitted_signal[:, i], h[i], mode='valid')
    return received_signal

# 信道估计和均衡
def channel_estimation_and_equalization(received_signal, h, fftsize, cpsize):
    equalized_signal = np.zeros((fftsize + cpsize, mimo_antennas), dtype=complex)
    for i in range(mimo_antennas):
        # 估计信道响应（这里简化为已知的信道响应）
        estimated_channel = np.fft.fftshift(np.fft.fft(received_signal[cpsize:cpsize+fftsize, i], fftsize))
        # 均衡
        equalized_signal[cpsize:cpsize+fftsize, i] = received_signal[cpsize:cpsize+fftsize, i] / estimated_channel
    return equalized_signal

# OFDM解调
def ofdm_demodulate(received_signal, fftsize, cpsize):
    ofdm_time = received_signal[cpsize:cpsize+fftsize]
    symbols = np.fft.fftshift(np.fft.fft(ofdm_time, fftsize))
    return symbols

# 主函数
def main():
    # 生成QAM符号
    symbols = generate_qam_symbols(databits)

    # 初始化transmitted_signal数组
    transmitted_signal = np.zeros((fftsize + cpsize, mimo_antennas), dtype=complex)

    # 生成随机信道
    h = np.random.randn(mimo_antennas, fftsize) + 1j * np.random.randn(mimo_antennas, fftsize)

    # OFDM调制
    for i in range(mimo_antennas):
        transmitted_signal[:, i] = ofdm_modulate(symbols, fftsize, cpsize)

    # 通过信道
    received_signal = channel_model(transmitted_signal, h)

    # 信道估计和均衡
    equalized_signal = channel_estimation_and_equalization(received_signal, h, fftsize, cpsize)

    # OFDM解调
    demodulated_symbols = np.zeros((N // 2, mimo_antennas), dtype=complex)
    for i in range(mimo_antennas):
        demodulated_symbols[:, i] = ofdm_demodulate(equalized_signal[:, i], fftsize, cpsize)

    print("发送端符号:\n", symbols)
    print("接收端解调符号:\n", demodulated_symbols)

if __name__ == "__main__":
    main()
